'''
https://leetcode.cn/problems/count-good-triplets-in-an-array/description/
'''
from typing import List

from sortedcontainers import SortedList

class IndexTree:
    def __init__(self, N):
        self.nums = [0] * (N+1)
        self.N = N
    def add(self, i, v):
        while i <= self.N:
            self.nums[i] += v
            i += i & -i
    def _sum(self, i):
        res = 0
        while i:
            res += self.nums[i]
            i -= i & -i
        return res

class Solution:
    def goodTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        n = len(nums1)
        ids = {}
        for i, num in enumerate(nums1):
            ids[num] = i

        bst = SortedList()
        res = 0
        for i, num in enumerate(nums2):
            id = ids[num]
            left_small = bst.bisect_left(id)
            all_big = n - 1 - id
            left_big = i - left_small
            right_big = all_big - left_big
            res += left_small * right_big
            bst.add(id)
        return res
    def goodTriplets2(self, nums1: List[int], nums2: List[int]) -> int:
        n = len(nums1)
        ids = {}
        for i, num in enumerate(nums1):
            ids[num] = i

        tree = IndexTree(n)
        res = 0
        for i, num in enumerate(nums2):
            id = ids[num]
            left_small = tree._sum(id + 1)
            all_big = n - 1 - id
            left_big = i - left_small
            right_big = all_big - left_big
            res += left_small * right_big
            tree.add(id + 1, 1)
        return res




nums1 = [4,0,1,3,2]
nums2 = [4,1,0,2,3]
print(Solution().goodTriplets(nums1, nums2))
print(Solution().goodTriplets2(nums1, nums2))

